# query_ddl.py
# usage: 根据question找到对应table创建ddl
# python ./src/query/query_ddl.py --query_text "请问2024年各受理渠道工单所占的百分比分别是多少？"

import os,sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
from config import project_config
import json
import argparse
from query.query_vector import retrieve_similar_texts
from query.query_reranker import query_reranker

def query_ddl(question: str,debug=None) -> str:
    filtered_docs = retrieve_similar_texts(question,debug, top_k=10)
    if len(filtered_docs) == 1 :
        try:
            # 解析 JSON 字符串，提取 DDL 字段
            matched_doc = json.loads(filtered_docs[0])
            table_define = matched_doc.get("DDL")
        except Exception as e:
            print(f"❌ 解析 JSON 失败: {e}")
            return None
    else :
        request_body = {
            "model": project_config.reranker_model,
            "input": {
                "query": question,
                "documents": filtered_docs
            }
        }
        table_define = query_reranker(request_body) 
    return table_define 

def main():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description="提示词大模型生成靠谱的SQL")
    parser.add_argument(
        '--query_text',
        type=str,
        required=True,
        help='需要检索的文本字符串。'
    )
    args = parser.parse_args()
    text_to_query = args.query_text

    # 将 silent 参数传递给函数
    ddl=query_ddl(text_to_query)
    print(f"ddl={ddl}")

# --- 程序入口 ---
if __name__ == '__main__':
    main()
